import numpy as np
import pandas as pd
import scipy.sparse as sp
from typing import List, Dict, Tuple
import cooler
import os
import re
from scipy.sparse import coo_matrix, load_npz
from collections import defaultdict
import gc

class ContactProcessor:
    def __init__(self, resolution: int = 50000, max_distance: int = 2000000):
        """
        初始化ContactProcessor工具类
        
        参数:
            resolution (int): 基因组分箱分辨率（默认50kb）
            max_distance (int): 最大交互距离（默认2Mb）
        """
        self.resolution = resolution
        self.max_distance = max_distance
        self.bin_dict = {}  # 存储染色体分箱信息
        self.num_bins = 0   # 总bin数

    def load_contact_file(self, file_path: str) -> pd.DataFrame:
        """
        加载contact文件
        
        参数:
            file_path (str): contact文件路径
        
        返回:
            pd.DataFrame: 包含chrom1, pos1, chrom2, pos2的DataFrame
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Contact file not found: {file_path}")
        
        # 加载文件，假设文件没有header，且列顺序为chrom1, pos1, chrom2, pos2
        df = pd.read_csv(file_path, sep='\t', header=None, 
                         names=['chrom1', 'pos1', 'chrom2', 'pos2'])
        return df

    def create_bin_dict(self, chrom_size_path: str):
        """
        创建染色体分箱字典
        
        参数:
            chrom_sizes (Dict[str, int]): 染色体名称到大小的映射
        """
        self.bin_dict = {}
        self.num_bins = 0
        if chrom_size_path is not None:
            chrom_sizes = pd.read_csv(chrom_size_path, sep='\t', index_col=0, header=None).squeeze(axis=1)
            for chrom, size in chrom_sizes.items():
                num_bins = (size // self.resolution) + 1
                self.bin_dict[chrom] = (self.num_bins, self.num_bins + num_bins)
                self.num_bins += num_bins

    def pos_to_bin(self, chrom: str, pos: int) -> int:
        """
        将基因组位置转换为bin索引
        
        参数:
            chrom (str): 染色体名称
            pos (int): 基因组位置
        
        返回:
            int: bin索引
        """
        if chrom not in self.bin_dict:
            raise ValueError(f"Chromosome {chrom} not found in bin dictionary")
        
        start_bin, end_bin = self.bin_dict[chrom]
        bin_idx = start_bin + (pos // self.resolution)
        if bin_idx >= end_bin:
            raise ValueError(f"Position {pos} is out of range for chromosome {chrom}")
        
        return bin_idx

    def process_contacts(self, df: pd.DataFrame) -> sp.coo_matrix:
        """
        处理contact文件，生成稀疏矩阵
        
        参数:
            df (pd.DataFrame): 包含chrom1, pos1, chrom2, pos2的DataFrame
        
        返回:
            scipy.sparse.coo_matrix: 稀疏接触矩阵
        """
        rows, cols, data = [], [], []
        
        for _, row in df.iterrows():
            try:
                bin_i = self.pos_to_bin(row['chrom1'], row['pos1'])
                bin_j = self.pos_to_bin(row['chrom2'], row['pos2'])
                
                # 过滤超出最大距离的接触
                if abs(bin_i - bin_j) * self.resolution > self.max_distance:
                    continue
                
                rows.append(bin_i)
                cols.append(bin_j)
                data.append(1)  # 每条数据表示一个相互作用，count默认为1
            except ValueError as e:
                print(f"Skipping invalid contact: {e}")
                continue
        
        return sp.coo_matrix((data, (rows, cols)), 
                            shape=(self.num_bins, self.num_bins))

    def normalize_matrix(self, matrix: sp.coo_matrix, method: str = 'ICE') -> sp.csr_matrix:
        """
        归一化接触矩阵
        
        参数:
            matrix (sp.coo_matrix): 原始稀疏矩阵
            method (str): 归一化方法（'ICE'或'KR'）
        
        返回:
            sp.csr_matrix: 归一化后的稀疏矩阵
        """
        if method == 'ICE':
            return cooler.balance_cooler(matrix, ignore_diags=2)
        elif method == 'KR':
            return cooler.balance_cooler(matrix, method='KR')
        else:
            raise ValueError(f"Unknown normalization method: {method}")

    def save_to_cool(self, matrix: sp.coo_matrix, output_path: str):
        """
        将稀疏矩阵保存为cool文件
        
        参数:
            matrix (sp.coo_matrix): 稀疏矩阵
            output_path (str): 输出文件路径
        """
        cooler.create_cooler(output_path, 
                            bins=pd.DataFrame({
                                'chrom': [k for k in self.bin_dict.keys()],
                                'start': [0] * len(self.bin_dict),
                                'end': [v[1] * self.resolution for v in self.bin_dict.values()],
                                'name': list(self.bin_dict.keys())
                            }), 
                            pixels=pd.DataFrame({
                                'bin1_id': matrix.row,
                                'bin2_id': matrix.col,
                                'count': matrix.data
                            }))

    def load_from_cool(self, cool_path: str) -> sp.coo_matrix:
        """
        从cool文件加载稀疏矩阵
        
        参数:
            cool_path (str): cool文件路径
        
        返回:
            sp.coo_matrix: 稀疏矩阵
        """
        c = cooler.Cooler(cool_path)
        return c.matrix(balance=False, as_pixels=True)

    def augment_data(self, matrix: sp.coo_matrix, mask_prob: float = 0.1) -> sp.coo_matrix:
        """
        数据增强：随机掩蔽接触点
        
        参数:
            matrix (sp.coo_matrix): 原始稀疏矩阵
            mask_prob (float): 掩蔽概率
        
        返回:
            sp.coo_matrix: 增强后的稀疏矩阵
        """
        mask = np.random.rand(matrix.nnz) > mask_prob
        return sp.coo_matrix((matrix.data[mask], 
                             (matrix.row[mask], matrix.col[mask])), 
                            shape=matrix.shape)
    
    def contact2npy(self, contact_path: str, chrom_sizes_path: str, output_path: str):
        """
        Convert contact data to chromosome-specific .npy matrices.
        
        Parameters:
            contact_path (str): Path to the contact data file.
            chrom_sizes_path (str): Path to the chromosome sizes file.
            output_path (str): Base directory to save output .npy files.
        """
        # 提取基础名称（如：GSM7682476_GM12878_cell_206）
        base_name = os.path.basename(contact_path).split('.')[0]
        # 创建输出目录（如：output_path/GSM7682476_GM12878_cell_206/）
        output_dir = os.path.join(output_path, base_name)
        os.makedirs(output_dir, exist_ok=True)

        data = pd.read_csv(
            contact_path,
            sep='\t', 
            usecols=[1,2,3,4], 
            names=['chr1', 'pos1', 'chr2', 'pos2'], 
            header=None)
        
        chrom_sizes = pd.read_csv(
            chrom_sizes_path, 
            sep='\t', 
            header=None, 
            usecols=[0, 1], 
            names=['chrom', 'size'])


        # 初始化接触矩阵字典
        contact_matrices = {}

        # 遍历每条染色体
        for chrom in chrom_sizes["chrom"].unique():
            # 获取当前染色体的长度
            chrom_size = chrom_sizes[chrom_sizes["chrom"] == chrom]["size"].values[0]
            
            # 计算矩阵大小
            matrix_size = (chrom_size // self.resolution) + 1
            
            # 初始化接触矩阵
            contact_matrix = np.zeros((matrix_size, matrix_size))
            
            # 筛选当前染色体的接触数据
            chrom_data = data[(data["chr1"] == chrom) & (data["chr2"] == chrom)]
            
            # 将位置映射到矩阵索引
            chrom_data["idx1"] = (chrom_data["pos1"] // self.resolution).astype(int)
            chrom_data["idx2"] = (chrom_data["pos2"] // self.resolution).astype(int)
            
            # 填充接触矩阵
            for _, row in chrom_data.iterrows():
                idx1 = row["idx1"]
                idx2 = row["idx2"]
                
                # 检查索引是否超出矩阵范围
                if idx1 >= matrix_size or idx2 >= matrix_size:
                    print(f"索引超出范围，跳过当前行。染色体: {chrom}, matrix_size: {matrix_size}, idx1: {idx1}, idx2: {idx2}, contact_file: {base_name}")
                    continue  # 跳过当前循环，进入下一个循环
                
                contact_matrix[idx1, idx2] += 1  # 累加计数
                contact_matrix[idx2, idx1] += 1  # 保持对称性
            
            # 保存接触矩阵
            contact_matrices[chrom] = contact_matrix

        # 保存接触矩阵
        for chrom, matrix in contact_matrices.items():
            filename = f"{base_name}_{chrom}.npy"
            filepath = os.path.join(output_dir, filename)
            np.save(filepath, matrix)

    def npy_to_cooler(self, npy_dir: str, chrom_sizes_path: str, output_cool_path: str):
        """
        将染色体特定的 .npy 接触矩阵转换为单一的 cooler 文件

        参数:
            npy_dir (str): 包含染色体特定 .npy 文件的目录
            chrom_sizes_path (str): 染色体大小文件的路径
            output_cool_path (str): 输出 cooler 文件的路径
        """
        # 1. 加载染色体大小
        chrom_sizes = pd.read_csv(chrom_sizes_path, sep='\t', header=None, 
                                names=['chrom', 'size'], dtype={'chrom': str})
        chromsizes = pd.Series(chrom_sizes.set_index('chrom')['size'].to_dict())

        # 2. 生成全局 bins
        bins = cooler.binnify(chromsizes, self.resolution)

        # 3. 初始化像素数据
        all_pixels = []

        # 4. 遍历处理每个染色体文件
        for filename in os.listdir(npy_dir):
            if not filename.endswith('.npy'):
                continue
                
            # 更准确地解析染色体名称（支持多种命名格式）
            chrom_match = re.search(r'[_\.]chr?([0-9XY]+)[_\.]', filename)
            if chrom_match:
                chrom_num = chrom_match.group(1)
                chrom = f"chr{chrom_num}" if not chrom_num.startswith("chr") else chrom_num
            else:
                # 尝试直接从文件名的最后部分提取
                chrom = os.path.splitext(filename)[0].split('_')[-1]
                if not chrom.startswith("chr"):
                    chrom = f"chr{chrom}"
            
            if chrom not in chromsizes.index:
                print(f"警告: 染色体 {chrom} 未在 chrom_sizes 中找到，跳过")
                continue

            print(f"处理染色体: {chrom}, 文件: {filename}")
            
            # 加载染色体特定矩阵
            matrix_path = os.path.join(npy_dir, filename)
            matrix = np.load(matrix_path)
            
            # 检查矩阵维度与染色体大小是否匹配
            expected_bins = int(np.ceil(chromsizes[chrom] / self.resolution))
            if matrix.shape[0] != expected_bins:
                print(f"警告: 矩阵维度 ({matrix.shape[0]}) 与预期的bin数量 ({expected_bins}) 不匹配")
                print(f"将使用实际的矩阵维度进行处理")
            
            # 计算bins范围
            chrom_bins = bins[bins['chrom'] == chrom]
            if len(chrom_bins) == 0:
                print(f"错误: 在bins列表中找不到染色体 {chrom}")
                continue
                
            bin_offset = chrom_bins['start'].iloc[0] // self.resolution
            
            # 转换为COO格式
            coo = coo_matrix(matrix)
            
            # 仅保留上三角部分（包括对角线）以避免重复记录
            mask = coo.row <= coo.col
            rows = coo.row[mask]
            cols = coo.col[mask]
            data = coo.data[mask]
            
            # 准备像素数据
            pixels_dict = {
                'bin1_id': [],
                'bin2_id': [],
                'count': []
            }
            
            # 计算全局bin索引
            for i, (r, c, v) in enumerate(zip(rows, cols, data)):
                if v == 0:  # 跳过零值
                    continue
                    
                # 获取对应的全局bin ID
                bin1_id = int(bin_offset + r)
                bin2_id = int(bin_offset + c)
                
                # 确保索引在有效范围内
                if bin1_id >= len(bins) or bin2_id >= len(bins):
                    continue
                    
                pixels_dict['bin1_id'].append(bin1_id)
                pixels_dict['bin2_id'].append(bin2_id)
                pixels_dict['count'].append(v)
            
            # 添加到全局像素列表
            if pixels_dict['bin1_id']:
                chrom_pixels = pd.DataFrame(pixels_dict)
                all_pixels.append(chrom_pixels)
                print(f"染色体 {chrom} 添加了 {len(chrom_pixels)} 个交互点")
            else:
                print(f"警告: 染色体 {chrom} 没有有效的交互数据")
        
        # 5. 合并所有像素数据
        if not all_pixels:
            raise ValueError("没有找到有效的交互数据，无法创建cooler文件")
            
        pixels = pd.concat(all_pixels, ignore_index=True)
        
        # 6. 排序并去重
        pixels = pixels.sort_values(['bin1_id', 'bin2_id'])
        pixels = pixels.groupby(['bin1_id', 'bin2_id'], sort=False)['count'].sum().reset_index()
        
        # 7. 确保bins包含bin_id列
        bins['bin_id'] = np.arange(len(bins))
        
        # 8. 创建cooler文件
        cooler.create_cooler(
            cool_uri=output_cool_path,
            bins=bins[['chrom', 'start', 'end']],
            pixels=pixels,
            metadata={'resolution': self.resolution},
            ordered=True
        )
        
        print(f"成功创建Cooler文件: {output_cool_path}")
        print(f"总交互数: {len(pixels):,}")
        print(f"覆盖染色体: {', '.join(sorted(set(bins.loc[pixels['bin1_id']]['chrom'])))}")
    
    def npz_to_cooler_single_chrom(self, npz_path: str, chrom_name: str, chrom_size: int, resolution: int, cool_path: str):
        """
        将单个染色体特定的 .npz 接触矩阵转换为 cooler 文件
        
        参数:
            npz_path (str): .npz 文件路径
            chrom_sizes_path (str): 染色体大小文件的路径
            output_cool_path (str): 输出 cooler 文件的路径
        """
        # 读取稀疏矩阵
        matrix = load_npz(npz_path).tocoo()

        # 构造 bins 表
        num_bins = int(np.ceil(chrom_size / resolution))
        bins = pd.DataFrame({
            "chrom": chrom_name,
            "start": np.arange(0, num_bins * resolution, resolution),
            "end": np.arange(resolution, (num_bins + 1) * resolution, resolution)
        })
        bins.loc[bins["end"] > chrom_size, "end"] = chrom_size  # 最后一个bin不要超过染色体长度

        # 构造 pixels 表
        pixels = pd.DataFrame({
            "bin1_id": matrix.row,
            "bin2_id": matrix.col,
            "count": matrix.data.astype(np.float32)
        })

        # 确保 bin1_id <= bin2_id（Hi-C 是对称的，但 cooler 要求上三角）
        b1 = pixels["bin1_id"].values
        b2 = pixels["bin2_id"].values
        pixels["bin1_id"] = np.minimum(b1, b2)
        pixels["bin2_id"] = np.maximum(b1, b2)

        # 合并重复（对称位置汇总）并排序
        pixels = (
            pixels.groupby(["bin1_id", "bin2_id"], sort=True)["count"].sum().reset_index()
        )
        pixels = pixels.sort_values(["bin1_id", "bin2_id"]).reset_index(drop=True)

        # 创建cool文件
        cooler.create_cooler(
            cool_uri=cool_path,
            bins=bins,
            pixels=pixels,
            ordered=True,
            metadata={
                'resolution': resolution,
                'bin-size': resolution
            },
            dtypes={'count': np.float32}
        )


    def npz_to_cooler(self, npz_path: str, chrom_sizes_path: str, output_cool_path: str, key: str = 'schicatt'):
        """
        将染色体特定的 .npz 接触矩阵转换为单一的 cooler 文件

        参数:
            npz_dir (str): 包含染色体特定 .npz 文件的目录
            chrom_sizes_path (str): 染色体大小文件的路径
            output_cool_path (str): 输出 cooler 文件的路径
        """
        # 1. 加载染色体大小（使用有序字典保持染色体顺序）
        chrom_sizes = pd.read_csv(chrom_sizes_path, sep='\t', header=None, 
                                names=['chrom', 'size'], dtype={'chrom': str})
        chromsizes = pd.Series(chrom_sizes.set_index('chrom')['size'].to_dict())

        # 2. 生成全局 bins（添加 bin_id 列）
        bins = cooler.binnify(chromsizes.sort_index(), self.resolution)
        bins['bin_id'] = np.arange(len(bins))  # 添加显式 bin_id 列

        # 3. 创建染色体偏移量映射表
        chrom_offsets = bins.groupby('chrom').first()['bin_id'].to_dict()

        # 4. 初始化像素数据容器（使用更高效的数组存储）
        bin1_ids, bin2_ids, counts = [], [], []

        # 5. 遍历处理每个染色体文件
        # for filename in os.listdir(npz_dir):
        if npz_path.endswith('.npz'):
            # 使用更稳健的染色体名称解析方法
            print(npz_path)
            chrom = os.path.splitext(npz_path)[0].split('_')[-1]
            print(chrom)
            if chrom not in chrom_offsets:
                raise ValueError(f"警告: 染色体 {chrom} 未在 chrom_sizes 中找到，跳过")

            # 加载染色体特定矩阵
            data = np.load(npz_path)
            matrix = data[key]
            
            # 转换为 COO 格式并过滤上三角（避免重复计数）
            coo = coo_matrix(matrix)
            mask = coo.row <= coo.col  # 保留上三角和对角线
            rows = coo.row[mask]
            cols = coo.col[mask]
            data = coo.data[mask]

            # 计算全局 bin ID
            offset = chrom_offsets[chrom]
            global_rows = rows + offset
            global_cols = cols + offset

            # 批量添加数据
            bin1_ids.append(global_rows)
            bin2_ids.append(global_cols)
            counts.append(data)

        # 6. 合并所有染色体数据
        pixels = pd.DataFrame({
            'bin1_id': np.concatenate(bin1_ids),
            'bin2_id': np.concatenate(bin2_ids),
            'count': np.concatenate(counts)
        })

        # 7. 创建排序索引并去重（解决可能的重复计数）
        pixels = pixels.sort_values(['bin1_id', 'bin2_id'])
        pixels = pixels.groupby(['bin1_id', 'bin2_id'], sort=False)['count'].sum().reset_index()

        # 8. 创建 cooler 文件（添加元数据）
        cooler.create_cooler(
            cool_uri=output_cool_path,
            bins=bins[['chrom', 'start', 'end']],
            pixels=pixels,
            metadata={'resolution': 50000},  # 根据实际分辨率调整
            dtypes={'count': np.int32},  # 根据实际数据类型调整
            ordered=True
        )

        print(f"成功创建 Cooler 文件: {output_cool_path}")
        print(f"总交互数: {pixels['count'].sum():,}")
        print(f"覆盖基因组范围: {', '.join(chromsizes.index)}")

    def merge_cell(self, input_dir: str, chrom_sizes_path: str, output_dir: str):

        """
        遍历 input_dir 下所有 cell 文件夹，将所有 cell 内相同染色体的 .npz 文件合并（对重复 (idx1, idx2) 累加）
        并以稀疏矩阵形式 .npz 文件输出到 output_dir。

        Parameters:
        -----------
        input_dir : str
            存放多个 cell 文件夹的目录，每个 cell 文件夹下有若干个染色体对应的 .npz 文件。
        chrom_sizes_path : str
            染色体大小文件的路径，每行格式例如: 
                chr1    248956422
                chr2    242193529
                ...
        output_dir : str
            用于保存合并后 .npz 文件的目录。
        """
        os.makedirs(output_dir, exist_ok=True)

        # 读取染色体大小
        chrom_sizes = {}
        with open(chrom_sizes_path, 'r') as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) < 2:
                    continue
                chrom, size_str = parts[:2]
                chrom_sizes[chrom] = int(size_str)

        # 用正则表达式匹配文件名中的染色体（假设文件名中包含 "_chrX.npz" 格式）
        chr_pattern = re.compile(r"_chr(\w+)\.npz$")

        # 全局字典：存储每条染色体合并后的 (i,j) -> 累加值
        chrom_data_dict = defaultdict(lambda: defaultdict(float))
        # 用于记录每条染色体需要的最大矩阵尺寸
        chrom_matrix_sizes = {}

        # 遍历 input_dir 下的所有 cell 文件夹
        for cell_folder in os.listdir(input_dir):
            cell_folder_path = os.path.join(input_dir, cell_folder)
            if not os.path.isdir(cell_folder_path):
                continue

            # 遍历当前 cell 文件夹下所有 .npz 文件
            for filename in os.listdir(cell_folder_path):
                if not filename.endswith('.npz'):
                    continue

                match = chr_pattern.search(filename)
                if not match:
                    print(f"{filename} 不符合格式，跳过...")
                    continue

                chrom = "chr" + match.group(1)  # 例如 group(1)="1" -> chrom="chr1"
                filepath = os.path.join(cell_folder_path, filename)

                # 载入稀疏矩阵并转换为 COO 格式以便遍历非零元素
                mat = sp.load_npz(filepath)
                mat_coo = mat.tocoo()

                # 累加到全局字典中
                for i, j, v in zip(mat_coo.row, mat_coo.col, mat_coo.data):
                    chrom_data_dict[chrom][(i, j)] += v

                # 根据染色体大小信息或矩阵本身的尺寸更新最大矩阵尺寸
                if chrom in chrom_sizes:
                    matrix_size = (chrom_sizes[chrom] // self.resolution) + 1
                else:
                    # matrix_size = mat.shape[0]
                    continue
                if chrom not in chrom_matrix_sizes:
                    chrom_matrix_sizes[chrom] = matrix_size
                else:
                    chrom_matrix_sizes[chrom] = max(chrom_matrix_sizes[chrom], matrix_size)

        # 构建并保存每条染色体合并后的稀疏矩阵
        for chrom, idx_dict in chrom_data_dict.items():
            matrix_size = chrom_matrix_sizes[chrom]
            print(f"正在为 {chrom} 构建大小为 {matrix_size}x{matrix_size} 的稀疏矩阵...")

            rows, cols, data = [], [], []
            for (i, j), val in idx_dict.items():
                # 只取索引在合理范围内的数值
                if i < matrix_size and j < matrix_size:
                    rows.append(i)
                    cols.append(j)
                    data.append(val)

            # 构建稀疏矩阵（COO 格式）
            merged_sparse = sp.coo_matrix((data, (rows, cols)), shape=(matrix_size, matrix_size))
            
            # 保存为 npz 文件
            output_npz_path = os.path.join(output_dir, f"merged_{chrom}.npz")
            sp.save_npz(output_npz_path, merged_sparse)
            print(f"已保存 {chrom} 到: {output_npz_path}")

    def contact2npz(self, contact_path: str, chrom_sizes_path: str, output_path: str):
        """
        Convert contact data to chromosome-specific sparse .npz matrices.
        
        Parameters:
            contact_path (str): Path to the contact data file.
            chrom_sizes_path (str): Path to the chromosome sizes file.
            output_path (str): Base directory to save output .npz files.
        """
        # Extract base name (e.g., GSM7682476_GM12878_cell_206)
        base_name = os.path.basename(contact_path).split('.')[0]
        # Create output directory (e.g., output_path/GSM7682476_GM12878_cell_206/)
        output_dir = os.path.join(output_path, base_name)
        os.makedirs(output_dir, exist_ok=True)

        # Load contact data
        data = pd.read_csv(
            contact_path,
            sep='\t',
            usecols=[1, 2, 3, 4],
            names=['chr1', 'pos1', 'chr2', 'pos2'],
            header=None
        )
        
        # Load chromosome sizes
        chrom_sizes = pd.read_csv(
            chrom_sizes_path,
            sep='\t',
            header=None,
            usecols=[0, 1],
            names=['chrom', 'size']
        )

        # Process each chromosome
        for chrom in chrom_sizes["chrom"].unique():
            # Get chromosome size
            chrom_size = chrom_sizes[chrom_sizes["chrom"] == chrom]["size"].values[0]
            
            # Calculate matrix size
            matrix_size = (chrom_size // self.resolution) + 1
            
            # Collect interactions in a list of (row, col, value) tuples
            interactions = []
            chrom_data = data[(data["chr1"] == chrom) & (data["chr2"] == chrom)]
            chrom_data["idx1"] = (chrom_data["pos1"] // self.resolution).astype(int)
            chrom_data["idx2"] = (chrom_data["pos2"] // self.resolution).astype(int)
            
            for _, row in chrom_data.iterrows():
                idx1 = row["idx1"]
                idx2 = row["idx2"]
                
                # Check if indices are within matrix bounds
                if idx1 >= matrix_size or idx2 >= matrix_size:
                    print(f"索引超出范围，跳过当前行。染色体: {chrom}, matrix_size: {matrix_size}, idx1: {idx1}, idx2: {idx2}, contact_file: {base_name}")
                    continue
                
                # Add interaction and its symmetric counterpart
                interactions.append((idx1, idx2, 1))
                if idx1 != idx2:
                    interactions.append((idx2, idx1, 1))
            
            # Create sparse matrix
            if interactions:
                rows, cols, values = zip(*interactions)
                contact_matrix = sp.coo_matrix((values, (rows, cols)), shape=(matrix_size, matrix_size)).tocsr()
            else:
                contact_matrix = sp.csr_matrix((matrix_size, matrix_size))
            
            # Save sparse matrix as .npz
            filename = f"{base_name}_{chrom}.npz"
            filepath = os.path.join(output_dir, filename)
            sp.save_npz(filepath, contact_matrix)
    
    def array_to_coo(array):
        """
        Convert a regular 2D NumPy array to a scipy.sparse.coo_matrix.

        Parameters:
        - array (numpy.ndarray): The input 2D array.

        Returns:
        - scipy.sparse.coo_matrix: The converted COO matrix.
        """
        # Find the non-zero elements in the array
        row, col = np.nonzero(array)

        # Get the values of the non-zero elements
        data = array[row, col]

        # Create the COO matrix
        coo_mat = sp.coo_matrix((data, (row, col)), shape=array.shape)

        return coo_mat
    
    def downsample_npz(self, input_filepath: str, proportions: list, output_path: str) -> None:
        """
        对指定路径的稀疏矩阵 npz 文件进行下采样操作，
        根据下采样比例列表生成对应的下采样矩阵，并保存至输出路径。
        
        参数:
            input_filepath (str): 需要下采样的 npz 文件路径。
            proportions (list of float): 下采样比例的列表，每个比例应在 (0, 1] 范围内。
            output_path (str): 输出文件保存的目录。
        """
        # 加载稀疏矩阵
        matrix = sp.load_npz(input_filepath)
        print(f"加载矩阵，原始非零元素数: {matrix.nnz}")
        
        # 获取文件基本名（不含扩展名）
        base_name = os.path.basename(input_filepath).replace('.npz', '')
        
        # 确保输出目录存在
        os.makedirs(output_path, exist_ok=True)
        
        # 针对每个下采样比例生成下采样矩阵并保存
        for p in proportions:
            # 复制矩阵，避免多次操作修改原矩阵
            matrix_copy = matrix.copy()
            # 对非零数据进行二项分布下采样（注意先转换为整数）
            original_data = matrix_copy.data.astype(int)
            matrix_copy.data = np.random.binomial(original_data, p)
            
            # 剔除采样后值为0的元素
            matrix_copy.eliminate_zeros()
            
            # 构造输出文件名，例如：原文件名_downsampled_50.npz
            output_filename = f"downsampled_{int(p * 100)}_{base_name}.npz"
            output_filepath = os.path.join(output_path, output_filename)
            
            # 保存下采样后的稀疏矩阵
            sp.save_npz(output_filepath, matrix_copy)
            print(f"下采样比例 {p}：保存文件 {output_filepath}，下采样后非零元素数: {matrix_copy.nnz}")

    def split_into_blocks(self, input_file: str, output_file: str):
        """
        将 input_dir 下所有合并后的 .npz 文件（格式为 (matrix_size, matrix_size)）转换为
        形状 (N, 1, 40, 40) 的块，并保存到 output_dir 中。
        
        Parameters:
        -----------
        input_dir : str
            包含合并后的 .npz 文件的目录。
        output_dir : str
            输出块文件的目录。
        """
        # os.makedirs(output_dir, exist_ok=True)
        k = 2000000 // self.resolution

        if not input_file.endswith('.npz'):
            raise ValueError(f'{input_file} File format error')
        
        # 解析染色体ID
        filename = os.path.basename(input_file)
        chr_pattern = re.compile(r"_chr([A-Za-z0-9]+)(_downsampled)?")
        match = chr_pattern.search(filename)
        # match = re.match(r'^chr([A-Za-z0-9]+)', filename)
        # print(match.group(1))
        # raise ValueError(f"文件名格式错误: 应以chr开头，如chr1.npz")
        if not match:
            raise ValueError(f"文件名格式错误: 应以chr开头，如chr1.npz")
        chrom_str = match.group(1)

        # 转换染色体标识符
        if chrom_str == 'X':
            chrom_id = 23
        elif chrom_str == 'Y':
            chrom_id = 24
        elif chrom_str == 'M':
            chrom_id = 25
        else:
            try:
                chrom_id = int(chrom_str)
            except ValueError:
                raise ValueError(f"无效的染色体标识符: {chrom_str}")


        
        # 处理单个文件
        merged_sparse = sp.load_npz(input_file)
        matrix_size = merged_sparse.shape[0]
        merged_coo = merged_sparse.tocoo()
        rows, cols, data = merged_coo.row, merged_coo.col, merged_coo.data
        
        # 计算每个元素的块索引
        block_dict = defaultdict(list)
        for i, j, val in zip(rows, cols, data):
            ib = i // k
            jb = j // k
            local_i = i % k
            local_j = j % k
            block_dict[(ib, jb)].append((local_i, local_j, val))
        
        # 确定块的数量
        num_blocks_per_dim = (matrix_size + k - 1) // k
        blocks = []
        inds_list = []

        # 生成所有块
        for ib in range(num_blocks_per_dim):
            for jb in range(num_blocks_per_dim):
                block = np.zeros((k, k), dtype=merged_coo.dtype)
                if (ib, jb) in block_dict:
                    for li, lj, v in block_dict[(ib, jb)]:
                        block[li, lj] = v
                # 调整形状为 (1, 1, 40, 40)
                block = block.reshape(1, 1, k, k)
                blocks.append(block)

                # 生成索引条目
                inds_entry = [
                    chrom_id,
                    k,        # 块大小
                    ib * k,  # 起始行坐标
                    jb * k,  # 起始列坐标
                ]
                inds_list.append(inds_entry)
        
        # 合并所有块
        if blocks:
            stacked_blocks = np.concatenate(blocks, axis=0)
        else:
            stacked_blocks = np.empty((0, 1, k, k), dtype=merged_coo.dtype)

        inds_array = np.array(inds_list, dtype=np.int32)
        # 保存结果
        np.savez_compressed(output_file, data=stacked_blocks, inds=inds_array)
        print(f"Saved {filename} data with shape {stacked_blocks.shape}")
        print(f"Saved {filename} inds with shape {inds_array.shape}")

    def npz_to_npy(self, npz_file: str, output_file: str = None, key: str = None):
        """
        将npz格式文件转换为npy格式文件
        
        参数:
            npz_file (str): npz文件路径
            output_file (str, optional): 输出npy文件路径，如果为None则使用原文件名替换扩展名
            key (str, optional): 要提取的数组键名，如果为None则提取第一个键
        
        返回:
            str: 保存的npy文件路径
        """
        # 如果未指定输出文件，则使用原文件名并更改扩展名
        if output_file is None:
            output_file = os.path.splitext(npz_file)[0] + '.npy'
        
        # 判断是否为scipy.sparse格式的npz文件
        try:
            # 尝试作为scipy.sparse稀疏矩阵加载
            sparse_matrix = sp.load_npz(npz_file)
            # 转换为密集数组
            dense_array = sparse_matrix.toarray()
            # 保存为npy文件
            np.save(output_file, dense_array)
            print(f"已将scipy.sparse稀疏矩阵npz文件转换为npy格式：{output_file}")
            print(f"数组形状：{dense_array.shape}")
            print(f"稀疏度：{sparse_matrix.nnz / (sparse_matrix.shape[0] * sparse_matrix.shape[1]):.6f}")
        except Exception as e:
            print(f"无法作为稀疏矩阵加载，尝试作为普通npz文件加载: {e}")
            # 加载普通npz文件
            data = np.load(npz_file)
            
            # 确定要提取的键
            if key is not None:
                if key not in data:
                    raise ValueError(f"键名 '{key}' 在npz文件中不存在")
                array = data[key]
            else:
                # 提取第一个键
                array = data[list(data.keys())[0]]
            
            # 保存为npy文件
            np.save(output_file, array)
            print(f"已将普通npz文件转换为npy格式：{output_file}")
            print(f"数组形状：{array.shape}")
        
        return output_file
        
    def npy_to_cool_single_chrom(self, npy_file: str, chrom: str = None, resolution: int = None, output_cool_path: str = None, chrom_sizes_path: str = None):
        """
        将单个染色体的npy文件转换为cool格式，保证染色体信息的正确映射
        
        参数:
            npy_file (str): npy文件路径
            chrom (str, optional): 染色体名称（例如'chr3'），如果为None则尝试从文件名解析
            resolution (int, optional): 分辨率（例如5000），如果为None则使用实例的分辨率
            output_cool_path (str, optional): 输出的cool文件路径，如果为None则使用npy文件名替换扩展名
            chrom_sizes_path (str, optional): 染色体大小文件路径，用于验证染色体信息
        
        返回:
            str: 创建的cool文件路径
        """
        # 设置默认值
        if resolution is None:
            resolution = self.resolution
        
        # 自动从文件名解析染色体名称
        if chrom is None:
            base_name = os.path.basename(npy_file)
            # 尝试多种命名模式
            chrom_match = re.search(r'[_\.]chr?([0-9XY]+)[_\.]', base_name)
            if chrom_match:
                chrom_num = chrom_match.group(1)
                chrom = f"chr{chrom_num}" if not chrom_num.startswith("chr") else chrom_num
            else:
                # 尝试直接从文件名的最后部分提取
                chrom = os.path.splitext(base_name)[0].split('_')[-1]
                if not chrom.startswith("chr"):
                    chrom = f"chr{chrom}"
            print(f"自动检测染色体名称: {chrom}")
        
        # 设置默认输出路径
        if output_cool_path is None:
            output_cool_path = os.path.splitext(npy_file)[0] + '.cool'
        
        # 验证染色体信息（如果提供了chrom_sizes_path）
        if chrom_sizes_path:
            chrom_sizes = pd.read_csv(chrom_sizes_path, sep='\t', header=None, 
                                    names=['chrom', 'size'], dtype={'chrom': str})
            if chrom not in chrom_sizes['chrom'].values:
                print(f"警告: 染色体 {chrom} 未在提供的染色体大小文件中找到")
            else:
                expected_size = chrom_sizes[chrom_sizes['chrom'] == chrom]['size'].values[0]
                print(f"染色体 {chrom} 大小: {expected_size}bp")
        
        # 加载npy文件
        matrix = np.load(npy_file)
        matrix_size = matrix.shape[0]
        print(f"矩阵维度: {matrix_size}x{matrix_size}")
        
        # 创建bins DataFrame
        bins = pd.DataFrame({
            'chrom': [chrom] * matrix_size,
            'start': np.arange(matrix_size) * resolution,
            'end': (np.arange(matrix_size) + 1) * resolution,
        })
        
        # 转换为COO格式 - 仅考虑非零元素
        rows, cols = np.nonzero(matrix)
        values = matrix[rows, cols]
        
        # 仅保留上三角部分（包括对角线）以避免重复记录
        mask = rows <= cols
        rows, cols, values = rows[mask], cols[mask], values[mask]
        
        # 创建pixels DataFrame
        pixels = pd.DataFrame({
            'bin1_id': rows,
            'bin2_id': cols,
            'count': values
        })
        
        # 创建cool文件
        cooler.create_cooler(
            cool_uri=output_cool_path,
            bins=bins,
            pixels=pixels,
            metadata={'resolution': resolution, 'chrom': chrom},
            ordered=True
        )
        
        print(f"成功创建单染色体Cooler文件: {output_cool_path}")
        print(f"染色体: {chrom}")
        print(f"分辨率: {resolution}bp")
        print(f"矩阵大小: {matrix_size}x{matrix_size}")
        print(f"总交互数: {len(values):,}")
        
        print("\n将cool文件转换为mcool的提示:")
        print("1. 使用cooler工具: cooler zoomify -r 1000,2000,5000,10000 your_file.cool")
        print("2. 或在Python中: import cooler; cooler.zoomify_cooler('your_file.cool', [1000,2000,5000,10000])")